rm(list=ls())

KFSurvey = function (num, y, mu0, Sigma0, Phi, Gam, cQ, cR, Rtdiag, input){
    Q = t(cQ) %*% cQ
    R = t(cR) %*% cR
    Phi = as.matrix(Phi)
    pdim = nrow(Phi)
    y = as.matrix(y)
    qdim = ncol(y)
    rdim = ncol(as.matrix(input))
    Ups = matrix(0, pdim, rdim)
    Ups = as.matrix(Ups)
    Gam = as.matrix(Gam) # qdim x rdim
    ut = matrix(input, num, rdim)
    Rtdiag=as.matrix(Rtdiag)
    xp = array(NA, dim = c(pdim, 1, num))
    Pp = array(NA, dim = c(pdim, pdim, num))
    xf = array(NA, dim = c(pdim, 1, num))
    Pf = array(NA, dim = c(pdim, pdim, num))
    innov = array(NA, dim = c(qdim, 1, num))
    sig = array(NA, dim = c(qdim, qdim, num))
    x00 = as.matrix(mu0, nrow = pdim, ncol = 1)
    P00 = as.matrix(Sigma0, nrow = pdim, ncol = pdim)
    xp[, , 1] = Phi %*% x00 + Ups %*% ut[1, ]
    Pp[, , 1] = Phi %*% P00 %*% t(Phi) + Q
    # no missing data, make B fixed
    B = matrix(1, nrow = qdim, ncol = pdim)
    Rt = R + diag(Rtdiag[1,], nrow=qdim, ncol=qdim)
    sig[,,1] = B %*% Pp[, , 1] %*% t(B) + Rt 
    siginv = chol2inv(chol(sig[, , 1]))
    K = Pp[, , 1] %*% t(B) %*% siginv
    innov[, , 1] = y[1, ] - B %*% xp[, , 1] - Gam %*% ut[1, ]
    xf[, , 1] = xp[, , 1] + K %*% innov[, , 1]
    Pf[, , 1] = Pp[, , 1] - K %*% B %*% Pp[, , 1]
    sigmat = as.matrix(sig[, , 1], nrow = qdim, ncol = qdim)
    like = log(det(sigmat)) + t(innov[, , 1]) %*% siginv %*% innov[, , 1]
    
    for (i in 2:num) {
        if (num < 2) 
            break
        xp[, , i] = Phi %*% xf[, , i - 1] + Ups %*% ut[i, ]
        Pp[, , i] = Phi %*% Pf[, , i - 1] %*% t(Phi) + Q 
        fullsig = tcrossprod(B %*% Pp[, , i],B) + R + diag(Rtdiag[i,], nrow=qdim, ncol=qdim)
        sig[,,i] = fullsig
        siginv = chol2inv(chol(sig[, , i]))
        K = Pp[, , i] %*% t(B) %*% siginv
        innov[, , i] = y[i, ] - B %*% xp[, , i] - Gam %*% ut[i,]
        xf[, , i] = xp[, , i] + K %*% innov[, , i]
        Pf[, , i] = Pp[, , i] - K %*% B %*% Pp[, , i]
        sigmat = matrix(sig[, , i], nrow = qdim, ncol = qdim) 
        like = like + log(det(sigmat)) + t(innov[, , i])%*%siginv %*% innov[, , i]
    }
    like = 0.5 * like
    list(xp = xp, Pp = Pp, xf = xf, Pf = Pf, like = like, innov = innov, 
        sig = sig, Kn = K)
}


library(foreign)

## first change commented out commands in 2_linregresults.do to create statespacedata.dta
statadata <- read.dta("statespacedata.dta")
summary(statadata)
nrow(statadata)

# remove 1952 (1st quarter 1952 missing)
statadata <- statadata[4:nrow(statadata),]
## Re-scale for use on 0-1 scale of macro
statadata$partyics <- statadata$partyics/100
statadata$partyapprove <- statadata$partyapprove/100


y <- matrix(statadata$adjusted_gallup,nc=1)
var <- matrix(statadata$macro,nc=1)
exog <- cbind(statadata$lagmacro, statadata$party, statadata$partyics, statadata$partyapprove)

# now define likelihood estimator

Linn <- function(para) {
    cQ <- matrix(para[5],1,1)
    cR <- matrix(para[6],1,1)
    kf <- KFSurvey(nrow(y), y, .5,.25, matrix(1,1,1), matrix(para[1:4],nrow=1), cQ, cR, var, exog)
    return(kf$like)
}
initpar <- c(.888490, -.024747, .007300, .035900, .2,.01616)

## test if likelihood function works
Linn(initpar)
test <- optim(initpar,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=TRUE,control=(list(trace=1,REPORT=1)))

stderrors <- sqrt(diag(solve(test$hessian)))

##
## Results in Table C4 (reverse sign in sigma epsilon, since it is squared in estimator)
##

rbind(test$par,stderrors)



KSSurvey = function (num, y, mu0, Sigma0, Phi, Gam, cQ, cR, Rtdiag, input){
    kf = KFSurvey(num, y, mu0, Sigma0, Phi, Gam, cQ, cR, Rtdiag, input)
    pdim = nrow(as.matrix(Phi))
    xs = array(NA, dim = c(pdim, 1, num))
    Ps = array(NA, dim = c(pdim, pdim, num))
    J = array(NA, dim = c(pdim, pdim, num))
    xs[, , num] = kf$xf[, , num]
    Ps[, , num] = kf$Pf[, , num]
    for (k in num:2) {
        J[, , k - 1] = (kf$Pf[, , k - 1] %*% t(Phi)) %*% solve(kf$Pp[, 
            , k])
        xs[, , k - 1] = kf$xf[, , k - 1] + J[, , k - 1] %*% (xs[, 
            , k] - kf$xp[, , k])
        Ps[, , k - 1] = kf$Pf[, , k - 1] + J[, , k - 1] %*% (Ps[, 
            , k] - kf$Pp[, , k]) %*% t(J[, , k - 1])
    }
    x00 = mu0
    P00 = Sigma0
    J0 = as.matrix((P00 %*% t(Phi)) %*% solve(kf$Pp[, , 1]), 
        nrow = pdim, ncol = pdim)
    x0n = as.matrix(x00 + J0 %*% (xs[, , 1] - kf$xp[, , 1]), 
        nrow = pdim, ncol = 1)
    P0n = P00 + J0 %*% (Ps[, , k] - kf$Pp[, , k]) %*% t(J0)
    list(xs = xs, Ps = Ps, x0n = x0n, P0n = P0n, J0 = J0, J = J, 
        xp = kf$xp, Pp = kf$Pp, xf = kf$xf, Pf = kf$Pf, like = kf$like, 
        Kn = kf$K)
}


cQ <- matrix(test$par[5],1,1)
cR <- matrix(test$par[6],1,1)
smooth <- KSSurvey(nrow(y), y, .5, .25, matrix(1,1,1), matrix(test$par[1:4],nrow=1), cQ, cR, var, exog)

ytilde <- (matrix(unlist(smooth$xs),nc=1) + test$par[2]*mean(statadata$party) + test$par[3]*mean(statadata$partyics) + test$par[4]*mean(statadata$partyapprove))/(1-test$par[1])

## compare with MSAR estimates after running 4_table2regressions.do
compare <- read.dta("4_conditionalupdatingmeasures.dta")

cov2cor(cov(cbind(compare$mu,smooth$xs)))
pdf("figureC3.pdf",height=4,width=4)
plot(compare$mu,smooth$xs,ylab="MSAR",xlab="State Space",main=expression(paste("Estimates of ",alpha[t]," correlate at .998")))
dev.off()


library(ggplot2)
library(RColorBrewer)

mypal <- brewer.pal(3,"Set1")
Quarter <- rep(seq(as.Date("1953/2/1"),as.Date("2012/11/1"),"3 months"),2)
eqmacro <- c(ytilde,compare$ytilde)
estimate <- c(rep("State Space",length(ytilde)),rep("MSAR",length(compare$ytilde)))
ggplotdata <- data.frame(Quarter,Macropartisanship=eqmacro,Estimate=estimate)
compareplot <- ggplot(ggplotdata,aes(Quarter,Macropartisanship))+ geom_line(aes(linetype=Estimate,colour=Estimate),size=1.2) +  scale_linetype_manual("Estimate",breaks=c("State Space","MSAR"),values=c(2,1)) + scale_colour_manual("Estimate",breaks=c("State Space","MSAR"),values=c(mypal[1],mypal[2])) + ylab("Equilibrium Macropartisanship") + theme_bw() +  theme(legend.position=c(.85,.85),legend.title=element_text(face="bold"))

pdf("figureC4.pdf",height=6,width=8)
compareplot
dev.off()

yhat <- matrix(unlist(smooth$xs),nc=1) + exog%*%matrix(test$par[1:4],nc=1)

## Compare RMSE

#SS
res <-  (y- yhat)*100
sqrt(mean(res^2))
#MSAR
res <-  (y- compare$ypred)*100
sqrt(mean(res^2))

#SS alpha R2
sum((ytilde-mean(y))^2)/sum((y-mean(y))^2)
#MSAR alpha R2
sum((compare$ytilde-mean(y))^2)/sum((y-mean(y))^2)

#SS R2
sum((yhat-mean(y))^2)/sum((y-mean(y))^2)
#MSAR R2
sum((compare$ypred-mean(y))^2)/sum((y-mean(y))^2)


## Now try a hetreg specification
inittest <- test

KFSurvey = function (num, y, mu0, Sigma0, Phi, Gam, hetreg, cR, Rtdiag, input){
    R = t(cR) %*% cR
    Phi = as.matrix(Phi)
    pdim = nrow(Phi)
    y = as.matrix(y)
    qdim = ncol(y)
    rdim = ncol(as.matrix(input))
    Ups = matrix(0, pdim, rdim)
    Ups = as.matrix(Ups)
    Gam = as.matrix(Gam) # qdim x rdim
    ut = matrix(input, num, rdim)
    Rtdiag=as.matrix(Rtdiag)
    xp = array(NA, dim = c(pdim, 1, num))
    Pp = array(NA, dim = c(pdim, pdim, num))
    xf = array(NA, dim = c(pdim, 1, num))
    Pf = array(NA, dim = c(pdim, pdim, num))
    innov = array(NA, dim = c(qdim, 1, num))
    sig = array(NA, dim = c(qdim, qdim, num))
    x00 = as.matrix(mu0, nrow = pdim, ncol = 1)
    P00 = as.matrix(Sigma0, nrow = pdim, ncol = pdim)
    xp[, , 1] = Phi %*% x00 + Ups %*% ut[1, ]
    Pp[, , 1] = Phi %*% P00 %*% t(Phi) + diag(hetreg[1,],nrow=pdim,ncol=pdim)
    # no missing data, make B fixed
    B = matrix(1, nrow = qdim, ncol = pdim)
    Rt = R + diag(Rtdiag[1,], nrow=qdim, ncol=qdim)
    sig[,,1] = B %*% Pp[, , 1] %*% t(B) + Rt 
    siginv = chol2inv(chol(sig[, , 1]))
    K = Pp[, , 1] %*% t(B) %*% siginv
    innov[, , 1] = y[1, ] - B %*% xp[, , 1] - Gam %*% ut[1, ]
    xf[, , 1] = xp[, , 1] + K %*% innov[, , 1]
    Pf[, , 1] = Pp[, , 1] - K %*% B %*% Pp[, , 1]
    sigmat = as.matrix(sig[, , 1], nrow = qdim, ncol = qdim)
    like = log(det(sigmat)) + t(innov[, , 1]) %*% siginv %*% innov[, , 1]
    
    for (i in 2:num) {
        if (num < 2) 
            break
        xp[, , i] = Phi %*% xf[, , i - 1] + Ups %*% ut[i, ]
        Pp[, , i] = Phi %*% Pf[, , i - 1] %*% t(Phi) + diag(hetreg[i,], nrow=pdim,ncol=pdim)
        fullsig = tcrossprod(B %*% Pp[, , i],B) + R + diag(Rtdiag[i,], nrow=qdim, ncol=qdim)
        sig[,,i] = fullsig
        siginv = chol2inv(chol(sig[, , i]))
        K = Pp[, , i] %*% t(B) %*% siginv
        innov[, , i] = y[i, ] - B %*% xp[, , i] - Gam %*% ut[i,]
        xf[, , i] = xp[, , i] + K %*% innov[, , i]
        Pf[, , i] = Pp[, , i] - K %*% B %*% Pp[, , i]
        sigmat = matrix(sig[, , i], nrow = qdim, ncol = qdim) 
        like = like + log(det(sigmat)) + t(innov[, , i])%*%siginv %*% innov[, , i]
    }
    like = 0.5 * like
    list(xp = xp, Pp = Pp, xf = xf, Pf = Pf, like = like, innov = innov, 
        sig = sig, Kn = K)
}


statadata$approveshock <- abs(abs(statadata$partyapprove) - abs(statadata$lagpartyapprove/100))*100
statadata$econshock <- abs(abs(statadata$partyics)- abs(statadata$lagpartyics/100))*100
transition <- c(1,abs(statadata$party[-1]-statadata$party[-240])/2)
statadata$approveshock <- ifelse(transition==1,0,statadata$approveshock)
statadata$econshock <- ifelse(transition==1,0,statadata$econshock)

## remove last 2 years
newstatadata <- statadata[1:232,]
y <- matrix(newstatadata$adjusted_gallup,nc=1)
var <- matrix(newstatadata$macro,nc=1)
exog <- cbind(newstatadata$lagmacro, newstatadata$party, newstatadata$partyics, newstatadata$partyapprove)


## first with only constant
hetexog <- cbind(rep(1,232))

Linn <- function(para) {
    hetreg <- exp(hetexog%*%matrix(c(para[5]),nc=1))
    cR <- matrix(para[6],1,1)
    kf <- KFSurvey(nrow(y), y, .5,.25, matrix(1,1,1), matrix(para[1:4],nrow=1), hetreg, cR, var, exog)
    return(kf$like)
}

initpar <- c(test$par[1:4],log(test$par[5]^2),test$par[6])

#undebug(KFSurvey)
#Linn(initpar)
test <- optim(initpar,Linn)
test <- optim(test$par,Linn,control=(list(trace=1,maxit=1000)))
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=TRUE,control=(list(trace=1,REPORT=1)))


stderrors <- sqrt(diag(solve(test$hessian)))
zscores <- test$par/stderrors
pvalues <- 1-pnorm(zscores)
rbind(test$par,stderrors,zscores,pvalues)

sqrt(exp(test$par[5]))
inittest$par[5]

hetexog <- cbind(rep(1,232),newstatadata$fpartyagendachange-mean(newstatadata$fpartyagendachange),newstatadata$fpartyideochange-mean(newstatadata$fpartyideochange))

Linn <- function(para) {
    hetreg <- exp(hetexog%*%matrix(c(para[5],para[6],para[7]),nc=1))
    cR <- matrix(para[8],1,1)
    kf <- KFSurvey(nrow(y), y, .5,.25, matrix(1,1,1), matrix(para[1:4],nrow=1), hetreg, cR, var, exog)
    return(kf$like)
}

initpar <- c(test$par[1:4],test$par[5],0,0,test$par[6])

#undebug(KFSurvey)
#Linn(initpar)
test <- optim(initpar,Linn)
test <- optim(test$par,Linn,control=(list(trace=1,maxit=1000)))
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=TRUE,control=(list(trace=1,REPORT=1)))
 

stderrors <- sqrt(diag(solve(test$hessian)))
zscores <- test$par/stderrors
pvalues <- 1-pnorm(zscores)
rbind(test$par,stderrors,zscores,pvalues)

hetexog <- cbind(hetexog,(newstatadata$approveshock - mean(newstatadata$approveshock)),(newstatadata$econshock - mean(newstatadata$econshock)))

Linn <- function(para) {
    hetreg <- exp(hetexog%*%matrix(c(para[5],para[6],para[7],para[8],para[9]),nc=1))
    cR <- matrix(para[10],1,1)
    kf <- KFSurvey(nrow(y), y, .5,.25, matrix(1,1,1), matrix(para[1:4],nrow=1), hetreg, cR, var, exog)
    return(kf$like)
}

initpar <- c(test$par[1:7],-.05,.09,-1*test$par[8])
test <- optim(initpar,Linn)
test <- optim(test$par,Linn,control=(list(trace=1,maxit=1000)))
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=FALSE,control=(list(trace=1,REPORT=1)))
test <- optim(test$par,Linn)
test <- optim(test$par,Linn,method="BFGS",hessian=TRUE,control=(list(trace=1,REPORT=1)))

stderrors <- sqrt(diag(solve(test$hessian)))
zscores <- test$par/stderrors
pvalues <- 1-pnorm(zscores)
rbind(test$par,stderrors,zscores,pvalues)

